#!/usr/bin/env python

import numpy, subprocess
import csv, timeit
import os
from scipy.io import FortranFile

def runcommand(cmd):
    tt0 = timeit.default_timer()
    retn = subprocess.call(cmd, shell=True)
    tt1 = timeit.default_timer()
    print('  {0:<40}  {1:3d}  Wall time: {2:.2f}'.format(cmd, retn, tt1-tt0))
    if retn !=0:
        exit()

def getdmab(myhf, nfrz):
    nocc = len(list(numpy.argwhere(myhf.mo_occ > 0).flatten()))
    nvir = len(list(numpy.argwhere(myhf.mo_occ == 0).flatten()))
    norb = nocc + nvir - nfrz
    norb0 = nocc + nvir

    print('nocc, nvir', nocc, nvir)

    f = FortranFile( 'cntomat.data', 'r' )
    a = f.read_reals( dtype='float64' )
    print('shape of a:', a.shape)
    dmab = a.reshape((norb,norb))
    print('shape of dmab:', dmab.shape)
    if nfrz == 0:
        numpy.savetxt('dmab.data', dmab)
    else:
        dmab1 = numpy.zeros((norb0, norb0), dtype = float)
        dmab1[nfrz:norb0, nfrz:norb0] = dmab[:,:]
        numpy.savetxt('dmab.data', dmab1)

def updatedropmo(filein):
    with open ('acted_orbital', 'r') as tmpinp:
         acted_orb = tmpinp.readlines()
         print(acted_orb)
    with open(filein,'r') as tmpinp:
         lists = tmpinp.readlines()
         iop_w = 0
    with open(filein,'w') as tmpinp:
         for i in range(0, len(lists)):
             if 'DROPMO' in lists[i]:
                 print('attention, dropmo can not be set in the prime ZMAT')
                 exit()
         for i in range(0, len(lists)):
             if 'CFOUR' in lists[i] or 'ACES2' in lists[i]:
                idx = lists[i].find(')')
                if idx>=0:
                   lists[i] = lists[i][:idx] + '\n'
                   acted_orb[0] = acted_orb[0].rstrip() + ')\n'
                lists.insert(i+1, acted_orb)
         for i in range(len(lists)):
             tmpinp.writelines((lists[i]))

def rmdropmo(filein):
    with open(filein, 'r') as tmpinp:
        lists = tmpinp.readlines()
    with open(filein, 'w') as tmpinp:
        ndel = []
        for i in range(len(lists)):
#        print(lists[i])
            idx = lists[i].find('DROPMO')
            if idx < 0:
               continue
            idx0 = lists[i].rfind(',', 0, idx)
            idxc = lists[i].find(',', idx)
            idxb = lists[i].find(')', idx)
            if idxc >=0 :
# case: DROPMO=..., .... 
               lists[i] = lists[i][:idx].rstrip() + lists[i][idxc+1:].lstrip()
            elif idx0 >= 0 and idxb < 0:
# case: ..., DROPMO=...
               lists[i] = lists[i][:idx0] + '\n'
            elif idx0 >= 0 and idxb >= 0:
# case: ..., DROPMO=...)
               lists[i] = lists[i][:idx0] + lists[i][idxb:]
            elif idx0 < 0 and idxb < 0:
# case: DROPMO=...
               ndel.append(i)
            else:
# case: DROPMO=...)
               ndel.append(i)
               lists[i-1] = lists[i-1].rstrip() + ')\n'
        for i in range(len(lists)):
            if i not in ndel:
               tmpinp.writelines((lists[i]))

def addpyscf(filein):
    with open(filein,'r') as tmpinp:
         lists = tmpinp.readlines()
    iop_w = 0
    with open(filein,'w') as tmpinp:
         for i in range(0, len(lists)):
             if '%readpyscf' in lists[i]:
               iop_w = 1
               break
         if iop_w == 0:
             lists.append('%readpyscf\n\n')
         for i in range(len(lists)):
             tmpinp.writelines((lists[i]))

def addnatorb(filein):
    with open(filein,'r') as tmpinp:
         lists = tmpinp.readlines()
    iop_w = 0
    with open(filein,'w') as tmpinp:
         for i in range(0, len(lists)):
             if '%natorb' in lists[i]:
               iop_w = 1
               break
         if iop_w == 0:
             lists.append('%natorb\n\n')
         for i in range(len(lists)):
             tmpinp.writelines((lists[i]))

def rmnatorb(filein):
    with open(filein,'r') as tmpinp:
         lists = tmpinp.readlines()
    with open(filein,'w') as tmpinp:
         for i in range(0, len(lists)):
             if '%natorb' in lists[i]:
               del lists[i]
               break
         for i in range(len(lists)):
             tmpinp.writelines((lists[i]))

def backupfile(filein):
    filenm = filein.rstrip()
    i = 0
    while True:
        fileout = filenm + '_' + str(i)
        if not os.path.exists(fileout):
           cmd = 'cp ' + filein + ' ' + fileout
           subprocess.call(cmd, shell=True)
           break
        i += 1
